# data = {'user' : [0, 1, 2], 'user_data' : {'1' : {'x' : [255, 245], 'y' : [0, 1]}, '2' : {'x' : [45, 54], 'y' : [1, 2]}}}
# x = data['user_data']['2']['x']
import copy
import json
import os
import random
import struct
import math
import numpy as np

def _data_norm(data):
    mean = np.mean(data)
    std_variance = math.sqrt(np.var(data))
    norm_data = (data - mean) / std_variance # data is a ndarray
    return norm_data

def data_norm(data):
    data_list = []
    for i in range(len(data)):
        norm_data_list = _data_norm(data[i]).tolist()  #每次对1 * 784大小的ndarray做标准化
        data_list.append(norm_data_list)
        
    return data_list

def unpickle(file):
    import pickle
    with open(file, 'rb') as fo:
        dict = pickle.load(fo, encoding='bytes')
    return dict

def extract_data(mode):
    y = []
    if mode == 'train':
        d1 = unpickle('/home/linjunke/downloads/cifar-100-python/train')
        total_x = d1[b'data']
        y = d1[b'coarse_labels']

    if mode == 'test':
        d = unpickle('/home/linjunke/downloads/cifar-100-python/test')
        total_x = d[b'data']
        y = d[b'coarse_labels']

    total_x = data_norm(total_x)
    return total_x, y

def _get_each_user_data(label_set, bin, x):

    each_user_data = {}
    data = []
    label = []
    #把数据和标签从箱子里拿出来

    for user_label in label_set:
        for i in bin[user_label]:
            label.append(user_label)
            data.append(x[i])

    #做一个同序shuffle，打乱顺序
    shuffle(data, label)
    #组装
    each_user_data['x'] = data
    each_user_data['y'] = label

    return each_user_data
#组装成最终结构
def get_noniid_data(user_num, user_label_mapping_dict, bin, x):
    data = {}
    total_user_data = {}
    data['users'] = [str(i) for i in range(user_num)] #注意是str（i），如果用int的话，在load_partition_data的时候会导致索引为int，
                                                    #然后报错keyerror，即索引错误。（因为转换成json格式之后，原来即便是int类型的key，也会变成str类型）
    
    for user, label_set in user_label_mapping_dict.items():
        total_user_data[str(user)] = _get_each_user_data(label_set, bin, x = x)

    data['user_data'] = total_user_data
    return data

def get_iid_data(user_num, data_num, x, y):
    data = []
    label = []
    batch_size = data_num // user_num
    for i in range(user_num):
        local_x = x[i * batch_size : (i + 1) * batch_size]
        local_y = y[i * batch_size : (i + 1) * batch_size]
        data.append(copy.deepcopy(local_x))
        label.append(copy.deepcopy(local_y))

    data_and_label = (data, label)
    return data_and_label

def shuffle(x, y):
    state = np.random.get_state()
    np.random.shuffle(x)
    np.random.set_state(state)
    np.random.shuffle(y)
    return x, y

def get_data(degree_of_iid, mode, x, y, user_label_dict, user_num):
    if mode == 'train':
        total_num = 50000
    else:
        total_num = 10000
    iid_num = int(degree_of_iid * total_num)
    noniid_num = total_num - iid_num
    iid_x = x[0 : iid_num]
    iid_y = y[0 : iid_num]
    noniid_x = x[iid_num : ]
    noniid_y = y[iid_num : ]

    bin_noniid = []
    for i in range(20):
        bin_noniid.append([])

    for i in range(noniid_num):
        label = noniid_y[i]
        bin_noniid[label].append(i)
    
    noniid_partial_dict = get_noniid_data(user_num = user_num, user_label_mapping_dict = user_label_dict, bin = bin_noniid, x = noniid_x)
    iid_partial_data_and_label = get_iid_data(user_num = user_num, data_num = iid_num, x = iid_x, y = iid_y)
    for i in range(user_num):
        noniid_partial_dict['user_data'][str(i)]['x'] += iid_partial_data_and_label[0][i]
        noniid_partial_dict['user_data'][str(i)]['y'] += iid_partial_data_and_label[1][i]
        shuffle(noniid_partial_dict['user_data'][str(i)]['x'], noniid_partial_dict['user_data'][str(i)]['y'])

    return noniid_partial_dict


train_x, train_y = extract_data(mode = 'train')
test_x, test_y = extract_data(mode = 'test')
#在data_norm()时已经做过tolist()了
# train_x = train_x.tolist()
# test_x = test_x.tolist()
user_label_dict = {0 : (1, 13, 3, 11), 1 : (19, 14, 7, 16), 2 : (9, 17, 4, 12), 3 : (2, 6, 15, 8), 4 : (10, 18, 5, 0)}
degree_of_iid = 0
user_num = 5

train_data_dict = get_data(degree_of_iid = degree_of_iid, mode = 'train', x = train_x, y = train_y, user_label_dict = user_label_dict, user_num = user_num)
#test_data的iid程度为1
test_data_dict = get_data(degree_of_iid = degree_of_iid, mode = 'test', x = test_x, y = test_y, user_label_dict = user_label_dict, user_num = user_num)

json.dump(train_data_dict, open("cifar100_train_100old.json", "w"))
# json.dump(test_data_dict, open("cifar100_test_0old_noniid.json", "w"))